import _thread as thread
import base64
import hashlib
import hmac
import json
from urllib.parse import urlparse, urlencode, unquote
import ssl
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
import websocket
from dotenv import load_dotenv
import os
from flask import Flask, request, Response, render_template, redirect, url_for, flash
from flask_cors import CORS
import queue
import sqlite3                                                                                                                                                        
# --- Flask App 初始化 ---
app = Flask(__name__)
CORS(app)
app.secret_key = os.urandom(24)

# --- 数据库代码 (保持不变) ---
DATABASE = 'chatbot.db'


def get_db_conn():
    conn = sqlite3.connect(DATABASE)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    with app.app_context():
        conn = get_db_conn()
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, role TEXT NOT NULL,
                content TEXT NOT NULL, scene TEXT NOT NULL, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )''')
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS products (
                id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL UNIQUE, category TEXT NOT NULL,
                description TEXT NOT NULL, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
            )''')
        conn.commit();
        conn.close();
        print("数据库已初始化 (messages & products 表).")


def save_message_to_db(session_id, role, content, scene):
    try:
        conn = get_db_conn();
        cursor = conn.cursor()
        cursor.execute('INSERT INTO messages (session_id, role, content, scene) VALUES (?, ?, ?, ?)',
                       (session_id, role, content, scene))
        conn.commit();
        conn.close()
    except Exception as e:
        print(f"数据库保存失败: {e}")


def get_products_by_category(category):
    conn = get_db_conn()
    products = conn.execute('SELECT name, description FROM products WHERE category = ?', (category,)).fetchall()
    conn.close()
    return products

# --- 新增数据库函数 ---
def get_all_products():
    """获取所有商品，按ID升序排序（最早添加的在最前面）"""
    conn = get_db_conn()
    # 核心改动：将排序方式改为 id ASC (ASC可以省略)
    products = conn.execute('SELECT * FROM products ORDER BY id').fetchall()
    conn.close()
    return products
def get_product_by_id(product_id):
    """根据ID获取单个商品"""
    conn = get_db_conn()
    product = conn.execute('SELECT * FROM products WHERE id = ?', (product_id,)).fetchone()
    conn.close()
    return product

# --- 加载 .env 和讯飞密钥 ---
load_dotenv()
spark_appid = os.getenv("APPID")
spark_api_key = os.getenv("API_KEY")
spark_api_secret = os.getenv("API_SECRET")


# --- 加载人设 (保持不变) ---
def load_prompts_from_json(file_path='prompts.json'):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            prompts = json.load(f)
            print("客服场景人设已从 prompts.json 成功加载。")
            return prompts
    except Exception as e:
        print(f"加载prompts.json失败: {e}")
        return {"general": "你是一个通用的AI助手。"}


SCENE_PROMPTS = load_prompts_from_json()


# --- 讯飞星火 WebSocket 相关代码 (流式版本，保持不变) ---
class Ws_Param:
    def __init__(self, APPID, APIKey, APISecret,
                 Spark_url): self.APPID = APPID; self.APIKey = APIKey; self.APISecret = APISecret; self.host = urlparse(
        Spark_url).netloc; self.path = urlparse(Spark_url).path; self.Spark_url = Spark_url

    def create_url(self): now = datetime.now(); date = format_date_time(mktime(
        now.timetuple())); sig_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"; sig_sha = hmac.new(
        self.APISecret.encode("utf-8"), sig_origin.encode("utf-8"),
        digestmod=hashlib.sha256).digest(); sig_sha_base64 = base64.b64encode(sig_sha).decode(
        "utf-8"); auth_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{sig_sha_base64}"'; auth = base64.b64encode(
        auth_origin.encode("utf-8")).decode("utf-8"); params = {"authorization": auth, "date": date,
                                                                "host": self.host}; return self.Spark_url + "?" + urlencode(
        params)


def on_message_stream(ws, message):
    data = json.loads(message)
    code = data["header"]["code"]
    if code != 0:
        error_msg = f"Spark请求错误: {data.get('header', {}).get('message', '未知错误')}"
        print(error_msg)
        ws.q.put(error_msg)
        ws.close()
        return
    choices = data["payload"]["choices"]
    status = choices["status"]
    content = choices["text"][0]["content"]
    if content:
        ws.q.put(content)
    if status == 2:
        ws.q.put(None)
        ws.close()


def spark_worker(appid, api_key, api_secret, Spark_url, domain, text_history, q):
    wsParam = Ws_Param(appid, api_key, api_secret, Spark_url)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message_stream, on_error=on_error, on_close=on_close,
                                on_open=on_open)
    ws.appid = appid;
    ws.domain = domain;
    ws.text_history = text_history;
    ws.q = q
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})


def call_spark_api_stream(appid, api_key, api_secret, Spark_url, domain, text_history):
    q = queue.Queue()
    thread.start_new_thread(spark_worker, (appid, api_key, api_secret, Spark_url, domain, text_history, q))
    while True:
        chunk = q.get()
        if chunk is None:
            break
        yield chunk


# --- Flask 路由 ---
@app.route('/')
def index(): return render_template('index.html')

# 1. 新增：知识管理主页
@app.route('/manage_knowledge')
def manage_knowledge():
    products = get_all_products()
    return render_template('manage_knowledge.html', products=products)

# 2. 修改：add_product 路由 (逻辑不变，但模板中会有返回链接)
@app.route('/add_product', methods=['GET', 'POST'])
def add_product():
    if request.method == 'POST':
        name=request.form['name']; category=request.form['category']; description=request.form['description']
        if not name or not category or not description: flash('所有字段都是必填项！', 'error')
        else:
            try:
                conn=get_db_conn(); conn.execute('INSERT INTO products (name, category, description) VALUES (?, ?, ?)',(name,category,description)); conn.commit(); conn.close()
                flash(f'商品 "{name}" 添加成功！', 'success')
                return redirect(url_for('manage_knowledge')) # 成功后返回管理主页
            except sqlite3.IntegrityError: flash(f'错误：商品名称 "{name}" 已存在！', 'error')
            except Exception as e: flash(f'添加失败: {e}', 'error')
        return redirect(url_for('add_product'))
    return render_template('add_product.html')

# 3. 新增：编辑商品路由
@app.route('/edit_product/<int:product_id>', methods=['GET', 'POST'])
def edit_product(product_id):
    product = get_product_by_id(product_id)
    if not product:
        flash('未找到该商品！', 'error')
        return redirect(url_for('manage_knowledge'))

    if request.method == 'POST':
        name = request.form['name']
        category = request.form['category']
        description = request.form['description']
        if not name or not category or not description:
            flash('所有字段都是必填项！', 'error')
        else:
            try:
                conn = get_db_conn()
                conn.execute('UPDATE products SET name = ?, category = ?, description = ? WHERE id = ?',
                             (name, category, description, product_id))
                conn.commit()
                conn.close()
                flash(f'商品 "{name}" 更新成功！', 'success')
                return redirect(url_for('manage_knowledge'))
            except sqlite3.IntegrityError:
                 flash(f'错误：商品名称 "{name}" 与其他商品冲突！', 'error')
            except Exception as e:
                flash(f'更新失败: {e}', 'error')
        # 如果更新失败，留在编辑页面
        return redirect(url_for('edit_product', product_id=product_id))

    return render_template('edit_product.html', product=product)

# 4. 新增：删除商品路由
@app.route('/delete_product/<int:product_id>', methods=['POST'])
def delete_product(product_id):
    product = get_product_by_id(product_id)
    if product:
        try:
            conn = get_db_conn()
            conn.execute('DELETE FROM products WHERE id = ?', (product_id,))
            conn.commit()
            conn.close()
            flash(f'商品 "{product["name"]}" 已被删除。', 'success')
        except Exception as e:
            flash(f'删除失败: {e}', 'error')
    else:
        flash('未找到该商品！', 'error')
    return redirect(url_for('manage_knowledge'))


@app.route('/stream_chat')
def stream_chat():
    scene = request.args.get('scene', 'general')
    session_id = request.args.get('session_id')
    history_str = request.args.get('history', '[]')
    history = json.loads(unquote(history_str))

    if not session_id:
        return Response("Session ID is required", status=400)

    user_message = history[-1]['content'] if history else ''
    if user_message:
        save_message_to_db(session_id, 'user', user_message, scene)

    knowledge_context = ""
    products = get_products_by_category(scene)
    if products:
        knowledge_context += "请严格根据以下提供的商品信息来回答问题(知识库上下文)：\n\n"
        for p in products:
            knowledge_context += f"--- 商品名称: {p['name']} ---\n商品详情: {p['description']}\n\n"
    base_prompt = SCENE_PROMPTS.get(scene, SCENE_PROMPTS['general'])
    enhanced_system_prompt = f"{base_prompt}\n\n{knowledge_context}"
    if not history or history[0].get('role') != 'system':
        messages_to_send = [{"role": "system", "content": enhanced_system_prompt}] + history
    else:
        history[0]['content'] = enhanced_system_prompt
        messages_to_send = history

    def generate():
        full_response = ""
        try:
            print("正在使用 Xunfei Spark API (Stream)...")
            stream = call_spark_api_stream(
                appid=spark_appid, api_key=spark_api_key, api_secret=spark_api_secret,
                Spark_url="wss://spark-api.xf-yun.com/v1/x1", domain="x1",
                text_history=messages_to_send
            )
            for chunk in stream:
                full_response += chunk
                yield f"data: {json.dumps({'content': chunk})}\n\n"
            yield "event: end\ndata: {}\n\n"
        except Exception as e:
            print(f"流式处理中发生错误: {e}")
            error_message = f"抱歉，处理时发生错误: {e}"
            yield f"data: {json.dumps({'content': error_message})}\n\n"
            yield "event: end\ndata: {}\n\n"
        if full_response:
            save_message_to_db(session_id, 'assistant', full_response, scene)

    return Response(generate(), mimetype='text/event-stream')


if __name__ == '__main__':
    init_db()


    def on_error(ws, error): print(f"### Spark WS Error: {error}")


    def on_close(ws, close_status_code, close_msg): pass


    def on_open(ws): thread.start_new_thread(run, (ws,))


    def run(ws): data = json.dumps(gen_params(ws.appid, ws.domain, ws.text_history)); ws.send(data)


    def gen_params(appid, domain, text_history): return {"header": {"app_id": appid, "uid": "1234"}, "parameter": {
        "chat": {"domain": domain, "temperature": 1.2, "max_tokens": 4096}},
                                                         "payload": {"message": {"text": text_history}}}


    app.run(host='0.0.0.0', port=5000, debug=True)